import sys
import os
import pathlib
import importlib
from enum import Enum

from PySide2 import QtWidgets
from PySide2 import QtCore
from PySide2 import QtGui
from PySide2 import QtUiTools
from maya import OpenMaya as om
from maya import OpenMayaUI as omui
from shiboken2 import wrapInstance
import maya.cmds as cmds

from .ui_resource import main_window
importlib.reload(main_window)

JOINT_TREE_ITEM_SIZE = QtCore.QSize(100, 30)



class JointTreeItem(QtWidgets.QWidget):

    On_Clock_Btn_clicked = QtCore.Signal()

    def __init__(self, name, parent=None):
        super().__init__(parent)

        self.name = name
        self._is_lock = False
        

        self.create_weidgtes()
        self.create_layouts()
        self.create_connections()

        self.init_state()
    
    def create_weidgtes(self):
        self.btn_lock = QtWidgets.QPushButton()
        self.btn_lock.setFixedSize(20, 20)

        self.btn_lock.setIcon(QtGui.QIcon(":Lock_OFF_grey.png"))

        icon = QtGui.QPixmap(":joint.svg").scaled(20, 20, QtCore.Qt.IgnoreAspectRatio, QtCore.Qt.SmoothTransformation)
        self.joint_label_icon = QtWidgets.QLabel()
        self.joint_label_icon.setPixmap(icon)

        self.joint_label_text = QtWidgets.QLabel(self.name)

    def create_layouts(self):
        main_hlayout = QtWidgets.QHBoxLayout(self)
        main_hlayout.addWidget(self.btn_lock)
        main_hlayout.addWidget(self.joint_label_icon)
        main_hlayout.addWidget(self.joint_label_text)
        main_hlayout.setContentsMargins(0,0,0,0)
        main_hlayout.setSpacing(5)
        main_hlayout.setStretch(0,1)
        main_hlayout.setStretch(1,1)
        main_hlayout.setStretch(2,15)

    def create_connections(self):
        self.btn_lock.clicked.connect(self.on_btn_clock_clicked)

    def init_state(self):
        if cmds.getAttr(f"{self.name}.liw"):
            self._is_lock = True
            self.btn_lock.setIcon(QtGui.QIcon(":Lock_ON.png"))
        else:
            self._is_lock = False
            self.btn_lock.setIcon(QtGui.QIcon(":Lock_OFF_grey.png"))

    def on_btn_clock_clicked(self):

        if self._is_lock:
            self._is_lock = False
            self.btn_lock.setIcon(QtGui.QIcon(":Lock_OFF_grey.png"))
            cmds.setAttr(f"{self.name}.liw", 0)
        else:
            self._is_lock = True
            self.btn_lock.setIcon(QtGui.QIcon(":Lock_ON.png"))
            cmds.setAttr(f"{self.name}.liw", 1)
        
        self.On_Clock_Btn_clicked.emit()


class CustomColorButton(QtWidgets.QWidget):
    
    color_changed = QtCore.Signal(QtGui.QColor)
    
    def __init__(self, parent=None):
        super().__init__(parent)
        
        self.create_control()
        
        self.set_size(50, 14)
        self.set_color(QtGui.QColor(255,255,255))
        
    def create_control(self):
        # 1. 创建 colorSliderGrp
        maya_window = cmds.window()
        self.__maya_colorslidergrp_name = cmds.colorSliderGrp()
        
        # 2. 找到 colorSliderGrp 控件
        color_slider_obj = omui.MQtUtil.findControl(self.__maya_colorslidergrp_name)
        if color_slider_obj:
            self.__color_slider_widget = wrapInstance(int(color_slider_obj), QtWidgets.QWidget)
            
        # 3. 将 colorSliderGrp 控件添加的当前 widget
        main_layout = QtWidgets.QVBoxLayout(self)
        main_layout.setObjectName("main_layout")
        main_layout.setContentsMargins(0,0,0,0)
        main_layout.addWidget(self.__color_slider_widget)
        
        # 4. 更新 colorSliderGrp
        self.__maya_colorslidergrp_name = omui.MQtUtil.fullName(int(color_slider_obj))
        
        # 5. 隐藏滑杆
        self.__slider_widget = self.__color_slider_widget.findChild(QtWidgets.QWidget, "slider")
        if self.__slider_widget:
            self.__slider_widget.hide()
            
        self.__color_widget = self.__color_slider_widget.findChild(QtWidgets.QWidget, "port")
        
        cmds.colorSliderGrp(self.__maya_colorslidergrp_name, e=True, changeCommand=self.on_color_changed)
        
        cmds.deleteUI(maya_window, window=True)
        
    def set_color(self, color):
        """
        设置颜色
        """
        color = QtGui.QColor(color)
        if color != self.get_color():
            cmds.colorSliderGrp(self.__maya_colorslidergrp_name, e=True, rgbValue=(color.redF(), color.greenF(), color.blueF()))
            self.on_color_changed()
        
    def get_color(self):
        """
        获取颜色
        """
        color = cmds.colorSliderGrp(self.__maya_colorslidergrp_name, q=True, rgbValue=True)
        
        return QtGui.QColor(color[0] * 255, color[1] * 255, color[2] * 255)
        
    def set_size(self, width, height):
        self.__color_slider_widget.setFixedWidth(width)
        self.__color_slider_widget.setFixedHeight(height)
        
    def on_color_changed(self, *args):
        self.color_changed.emit(self.get_color())


class SkinType(Enum):
    xsCustomSkin = "xsCustomSkin"
    skinCluster = "skinCluster"


class XS_SkinPainter(QtWidgets.QWidget, main_window.Ui_Form):

    def __init__(self):
        super().__init__()

        self.setupUi(self)

        self.setWindowFlags(QtCore.Qt.WindowStaysOnTopHint)
        self.setWindowTitle("XS SkinPainter")

        self.set_jointview_item_visible(1, False)

        # attrs
        self._skin = om.MFnDependencyNode()
        self._dagpath_skin_joints = []
        self._dagpath_str_skin_joints = []
        self._script_job_selection_changed = None
        self._script_job = None
        self._skin_type = None
        self._draw_mesh = ""

        self.init_widgets()

        self.create_connection()

    def create_connection(self):
        self.pushButton_paint.clicked.connect(self.initialize_selected_skined_mesh)
        self.radioButton_Replace.toggled.connect(self.on_radiobtn_toggled)
        self.radioButton_Add.toggled.connect(self.on_radiobtn_toggled)
        self.radioButton_Smooth.toggled.connect(self.on_radiobtn_toggled)

        self.dSpinBox_Opacity.valueChanged.connect(self.on_spinbox_opacity_chaned)
        self.dSpinBox_Value.valueChanged.connect(self.on_spinbox_value_chaned)

        self.hSlider_Opacity.valueChanged.connect(self.on_slider_opacity_changed)
        self.hSlider_Value.valueChanged.connect(self.on_slider_value_changed)

        self.treeWidget_Joints.itemClicked.connect(self.on_joint_tree_item_clicked)

    def init_widgets(self):
        """
        初始化设置
        """
        self.treeWidget_Joints.setHeaderHidden(True)
        self.treeWidget_Joints.setIconSize(JOINT_TREE_ITEM_SIZE)
        self.treeWidget_Joints.viewport().installEventFilter(self)

        self.colorP = QtWidgets.QWidget()
        self.colorP.setMinimumSize(50, 20)
        self.brushColorBtn = CustomColorButton(self.colorP)
        self.brushColorBtn.set_color(QtGui.QColor(255.0, 0.0, 0.0))
        self.brushColorBtn.color_changed.connect(self.on_brush_color_changed)
        self.hLayout_ColorBtn.addWidget(self.colorP)

        self.dSpinBox_Opacity.setValue(1.0)
        self.dSpinBox_Value.setValue(1.0)

    def eventFilter(self, watched:QtWidgets.QWidget, event:QtCore.QEvent):
        
        if watched == self.treeWidget_Joints.viewport():
            if event.type() == QtCore.QEvent.MouseButtonPress:
                pos = event.pos()
                item = self.treeWidget_Joints.itemAt(pos)
                if not item:
                    self.treeWidget_Joints.clearSelection()
                    self.treeWidget_Joints.clearFocus()
        
        return super().eventFilter(watched, event)

    def set_jointview_item_visible(self, index, visible):
        """
        设置joint面板显示隐藏
        """

        layout_item = self.vLayout_JointView.itemAt(index)
        if layout_item:
            layout_item.widget().setVisible(visible)

    def find_skin_with_selected_mesh(self):
        """
        在选择的模型上查找蒙皮节点
        """
        selected_objs = cmds.ls(selection=True)
        if not selected_objs:
            return
        
        selected_mesh = selected_objs[0]
        obj_relatives = cmds.listRelatives(selected_mesh, children=True, shapes=True)
        if not obj_relatives:
            om.MGlobal.displayError("Selected object must be a polymesh")
            return
        
        mesh_shape = obj_relatives[0]
        if cmds.nodeType(mesh_shape) != "mesh":
            om.MGlobal.displayError("Selected object must be a polymesh")
            return
        
        # find skin node
        self._skin = None
        for node in cmds.listHistory(mesh_shape):
            if cmds.nodeType(node) == SkinType.xsCustomSkin.value:
                self._skin = node
                self._skin_type = SkinType.xsCustomSkin
            elif cmds.nodeType(node) == SkinType.skinCluster.value:
                self._skin = node
                self._skin_type = SkinType.skinCluster

        return selected_mesh, self._skin

    def get_skin_joints(self):
        """
        返回所有参与蒙皮的骨骼
        """
        if not self._skin:
            return []
        
        msel_list = om.MSelectionList()
        msel_list.add(self._skin)

        mobj_skin = om.MObject()
        msel_list.getDependNode(0, mobj_skin)

        mfnskin_node:om.MFnDependencyNode = om.MFnDependencyNode(mobj_skin)

        if self._skin_type == SkinType.xsCustomSkin:
            plug_matrix:om.MPlug = mfnskin_node.findPlug("influenceMatrices", False)
        elif self._skin_type == SkinType.skinCluster:
            plug_matrix:om.MPlug = mfnskin_node.findPlug("matrix", False)

        numjoint = plug_matrix.numElements()
        
        msel_list = om.MSelectionList()
        mdagpath_array_joints = []
        for index in range(numjoint):
            plug_matrix_i:om.MPlug = plug_matrix.elementByLogicalIndex(index)
            mobj = plug_matrix_i.source().node()

            mfndag_node = om.MFnDagNode(mobj)
            mdagpath = om.MDagPath()
            mfndag_node.getPath(mdagpath)

            mdagpath_array_joints.append((index, mdagpath))

        # (index, dagpath_joint)
        self._dagpath_skin_joints = mdagpath_array_joints
        self._dagpath_str_skin_joints = []
        for index, dagpath in self._dagpath_skin_joints:
            self._dagpath_str_skin_joints.append(dagpath.fullPathName())

        return self._dagpath_skin_joints

    def initialize_selected_skined_mesh(self):
        """
        初始化选择的被蒙皮的模型
        """
        drawmesh, skin_node = self.find_skin_with_selected_mesh()
        if not skin_node:
            om.MGlobal.displayError("Selected mesh has no skin")
            return

        self.refrush_jointtree()
        self.set_jointview_item_visible(0, False)
        self.set_jointview_item_visible(1, True)

        # set tool
        self._draw_mesh = drawmesh
        self._brush_context = cmds.xsSkinBrush()
        cmds.setToolTo(self._brush_context)
        cmds.xsSkinBrush(self._brush_context, e=True, mesh=self._draw_mesh, updateInfluenceJoint=True)

        # init tool
        opacity = self.dSpinBox_Opacity.value()
        cmds.xsSkinBrush(self._brush_context, e=True, opacity=opacity)

        value = self.dSpinBox_Value.value()
        cmds.xsSkinBrush(self._brush_context, e=True, value=value)

        op = "add"
        if self.radioButton_Replace.isChecked():
            op = "replace"
        elif self.radioButton_Smooth.isChecked():
            op = "smooth"
        cmds.xsSkinBrush(self._brush_context, e=True, paintOperation=op)

        color = self.brushColorBtn.get_color()
        cmds.xsSkinBrush(self._brush_context, e=True, brushColor=(color.red()/255.0, color.green()/255.0, color.blue()/255.0))
        cmds.xsSkinBrush(self._brush_context, e=True, jointIndex=0)
        cmds.xsSkinBrush(self._brush_context, e=True, mesh=self._draw_mesh, updateInfluenceJoint=True)

        # create scriptJob
        self._script_job = cmds.scriptJob(event=("PostToolChanged", self.on_post_tool_changed))

    def closeEvent(self, event):
        self.clear_script_job()

    def refrush_jointtree(self):
        """
        刷新骨骼树
        """
        self.treeWidget_Joints.clear()

        dagpth_skin_joints = self.get_skin_joints()
        self.create_tree_item_recursive(dagpth_skin_joints)

    def create_tree_item_recursive(self, dagpath_array_jnts, depth=1):
        """
        递归创建骨骼树
        """
        if len(dagpath_array_jnts) == 0:
            return

        next_dagpath_array_jnts = []
        next_depth = depth + 1
        for item in dagpath_array_jnts:
            index = item[0]
            dagpath_jnt = item[1]

            current_depth = dagpath_jnt.length()
            if current_depth == depth:
                new_item = QtWidgets.QTreeWidgetItem()
                new_item.setData(0, QtCore.Qt.UserRole, (current_depth, dagpath_jnt.partialPathName(), dagpath_jnt.fullPathName(), index))
                new_item.setSizeHint(0, JOINT_TREE_ITEM_SIZE)

                parent_item:QtWidgets.QTreeWidgetItem = self.find_parent_treeitem(new_item)
                
                if parent_item:
                    parent_item.addChild(new_item)
                else:
                    self.treeWidget_Joints.addTopLevelItem(new_item)


                joint_item = JointTreeItem(new_item.data(0, QtCore.Qt.UserRole)[1])
                joint_item.On_Clock_Btn_clicked.connect(self.on_joint_item_clock_clicked)
                self.treeWidget_Joints.setItemWidget(new_item, 0, joint_item)
            else:
                next_dagpath_array_jnts.append((index, dagpath_jnt))

        self.create_tree_item_recursive(next_dagpath_array_jnts, next_depth)
        
    def dagpath_str_array_filter(self, dagpath_str_array):
        """
        过滤 dagptah
        1. 过滤非骨骼节点
        2. 过滤不参与蒙皮的骨骼节点
        """
        temp_array = []
        msel_list = om.MSelectionList()
        index = 0
        for item in dagpath_str_array:
            # 过滤非骨骼节点
            if cmds.nodeType(item) == "joint":
                # 过滤不参与蒙皮的骨骼节点
                msel_list.add(item)
                dagpath_item = om.MDagPath()
                msel_list.getDagPath(index, dagpath_item)
                item_full_path = dagpath_item.fullPathName()
                index += 1
                if item_full_path in self._dagpath_str_skin_joints:
                    temp_array.append(item)

        return temp_array

    def find_parent_treeitem(self, item:QtWidgets.QTreeWidgetItem):
        """
        返回item的父item
        """
        parent_item = None
        full_path_name = item.data(0, QtCore.Qt.UserRole)[2]
        dagpath_str_array = []
        for path_item in full_path_name.split("|"):
            dagpath_str_array.append(path_item)

        dagpath_str_array = dagpath_str_array[1:]
        dagpath_str_array = self.dagpath_str_array_filter(dagpath_str_array)
        # 路径对象个数为 1，说明为根节点
        if len(dagpath_str_array) == 1:
            return parent_item
        
        root_count = self.treeWidget_Joints.topLevelItemCount()
        root_item = None
        for i in range(root_count):
            temp_root:QtWidgets.QTreeWidgetItem = self.treeWidget_Joints.topLevelItem(i)
            if temp_root.data(0, QtCore.Qt.UserRole)[1] == dagpath_str_array[0]:
                root_item = temp_root
                break

        dagpath_str_array.pop(0)
        parent_item = self.find_parent_treeitem_recursive(root_item, dagpath_str_array)

        return parent_item

    def find_parent_treeitem_recursive(self, root_item:QtWidgets.QTreeWidgetItem, dagpath_str_array):
        """
        递归查找 item 的 父 item
        """
        if len(dagpath_str_array) == 1:
            return root_item
        
        item_count = root_item.childCount()
        next_root_item = None
        for i in range(item_count):
            temp_item:QtWidgets.QTreeWidgetItem = root_item.child(i)
            if temp_item.data(0, QtCore.Qt.UserRole)[1] == dagpath_str_array[0]:
                next_root_item = temp_item

        dagpath_str_array.pop(0)
        parent_item = self.find_parent_treeitem_recursive(next_root_item, dagpath_str_array)
        return parent_item

    def on_brush_color_changed(self, color:QtGui.QColor):
        brush_color = (color.redF(), color.greenF(), color.blueF())
        cmds.xsSkinBrush(self._brush_context, e=True, brushColor=(brush_color[0], brush_color[1], brush_color[2]))

    def clear_script_job(self):
        if self._script_job:
            cmds.evalDeferred(f"if cmds.scriptJob(exists={self._script_job}):\n\tcmds.scriptJob(kill={self._script_job}, force=True)")
            self._script_job = None

    def on_post_tool_changed(self):
        self.set_jointview_item_visible(0, True)
        self.set_jointview_item_visible(1, False)

        self.clear_script_job()

    def on_radiobtn_toggled(self, checked):
        if self.radioButton_Smooth.isChecked():
            self.dSpinBox_Value.setEnabled(False)
            self.hSlider_Value.setEnabled(False)
        else:
            self.dSpinBox_Value.setEnabled(True)
            self.hSlider_Value.setEnabled(True)

        if self.radioButton_Add.isChecked():
            cmds.xsSkinBrush(self._brush_context, e=True, paintOperation="add")
        elif self.radioButton_Replace.isChecked():
            cmds.xsSkinBrush(self._brush_context, e=True, paintOperation="replace")
        elif self.radioButton_Smooth.isChecked():
            cmds.xsSkinBrush(self._brush_context, e=True, paintOperation="smooth")

    def on_spinbox_opacity_chaned(self, val):
        self.hSlider_Opacity.setValue(int(val * 100))
        cmds.xsSkinBrush(self._brush_context, e=True, opacity=val)

    def on_spinbox_value_chaned(self, val):
        self.hSlider_Value.setValue(int(val * 100))

    def on_slider_opacity_changed(self, val):
        self.dSpinBox_Opacity.setValue(float(val * 0.01))
        cmds.xsSkinBrush(self._brush_context, e=True, opacity=float(val * 0.01))

    def on_slider_value_changed(self, val):
        self.dSpinBox_Value.setValue(float(val * 0.01))
        cmds.xsSkinBrush(self._brush_context, e=True, value=float(val * 0.01))

    def on_joint_tree_item_clicked(self, item:QtWidgets.QTreeWidgetItem, column:int):
        jnt_index = item.data(0, QtCore.Qt.UserRole)[3]
        jnt_index = int(jnt_index)
        cmds.xsSkinBrush(self._brush_context, e=True, jointIndex=jnt_index)
        cmds.select(clear=True)
        cmds.select(item.data(0, QtCore.Qt.UserRole)[1])

    def on_joint_item_clock_clicked(self):
        cmds.xsSkinBrush(self._brush_context, e=True, mesh=self._draw_mesh, updateInfluenceJoint=True)



if __name__ == "__main__":
    painter_wnd = XS_SkinPainter()
    painter_wnd.show()
